"""
LangChain RAG（检索增强生成）代理系统学习示例
===================================================
这个示例演示如何构建一个能够检索外部知识并回答问题的智能代理系统。
系统从网页加载知识，构建向量数据库，并创建能够检索信息的工具。

核心组件：
1. Document Loader: 文档加载器，从网页加载内容
2. Text Splitter: 文本分割器，将长文档切分为小块
3. Embeddings: 向量化模型，将文本转换为向量表示
4. Vector Store: 向量数据库，存储和检索向量化文档
5. Retriever Tool: 检索工具，让Agent能够搜索知识库
6. Tool Calling Agent: 工具调用代理，理解问题并选择合适工具

工作流程：
网页内容 → 文档加载 → 文本分割 → 向量化 → 向量存储 → 检索工具 → 智能代理
"""
import os

from dotenv import load_dotenv
from langchain import hub
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import create_retriever_tool
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 自动加载 .env 文件中的环境变量（用于API密钥等配置）
load_dotenv()

# ==================== 定义大语言模型 ====================
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen3-Coder-480B-A35B-Instruct", streaming=True, temperature=0.7) # 代码专用模型

# ==================== 知识库构建流程 ====================

# 步骤1：文档加载 - 从网页加载知识内容
print("正在加载网页知识内容...")
loader = WebBaseLoader("https://baike.baidu.com/item/%E4%BA%BA%E5%B7%A5%E6%99%BA%E8%83%BD/9180")
# 加载文档，返回Document对象列表，每个Document包含页面内容和元数据
docs = loader.load()
print(f"成功加载 {len(docs)} 个文档")

# 步骤2：文档切分 - 将长文档分割为适合处理的文本块
print("正在分割文档...")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,   # 每个文本块的最大字符数（约500-1000汉字）
    chunk_overlap=200  # 块之间的重叠字符数，保持上下文连贯性，避免信息割裂
)
documents = text_splitter.split_documents(docs)  # 执行切分操作
print(f"文档分割完成，共 {len(documents)} 个文本块")

# 步骤3：向量化处理 - 将文本转换为数值向量
print("正在初始化向量化模型...")
# 选项1：使用在线模型（需要网络）
# embeddings = HuggingFaceEmbeddings(
#     model_name="moka-ai/m3e-base",
#     model_kwargs={'device': 'cpu'}  # 使用CPU进行计算
# )

# 选项2：使用本地模型（推荐，更快更稳定）
embeddings = HuggingFaceEmbeddings(model_name="D:/models/BAAIbge-base-zh-v1.5", model_kwargs={'device': 'cpu'})

# 步骤4：创建向量数据库 - 存储文档向量用于相似度检索
print("正在构建向量数据库...")
vector = FAISS.from_documents(
    documents,    # 分割后的文档列表
    embeddings    # 向量化模型
)
print("向量数据库构建完成")

# 步骤5：创建检索器 - 用于在向量数据库中搜索相关文档
retriever = vector.as_retriever(
    search_type="similarity",  # 使用相似度搜索
    search_kwargs={"k": 3}     # 返回最相关的3个文档片段
)

# ==================== 工具创建 ====================

# 创建检索工具 - 让Agent能够搜索百度百科知识库
retriever_tool = create_retriever_tool(
    retriever,        # 检索器对象
    "baike_search",   # 工具名称（Agent调用时使用的标识）
    "搜索百度百科中关于人工智能的相关知识，当需要查询人工智能的定义、历史、应用等信息时使用此工具"  # 工具描述，Agent根据描述决定何时使用该工具
)

# 工具列表 - 可以添加多个工具，Agent会根据问题自动选择
tools = [retriever_tool]
print(f"已创建 {len(tools)} 个工具")

# ==================== 代理系统构建 ====================

# 获取预定义的提示词模板
print("正在加载提示词模板...")
prompt = hub.pull("hwchase17/openai-functions-agent")  # 从LangChain Hub拉取标准模板
# 这个模板包含了Agent的思考框架、工具使用说明等

# 创建工具调用代理
print("正在创建智能代理...")
agent = create_tool_calling_agent(
    llm,      # 大语言模型，负责理解和决策
    tools,    # 可用工具列表
    prompt    # 提示词模板，指导Agent的行为
)

# 创建代理执行器 - 管理Agent的执行流程
agent_executor = AgentExecutor(
    agent=agent,     # 代理实例
    tools=tools,     # 工具列表
    verbose=True,    # 显示详细执行过程（学习时建议开启）
    handle_parsing_errors=True  # 自动处理解析错误
)

# ==================== 测试代理系统 ====================

print("\n" + "="*50)
print("开始测试代理系统")
print("="*50)

# 测试1：基于知识库的问题（使用检索工具）
print("\n--- 测试1：知识检索问题 ---")
response1 = agent_executor.invoke({
    "input": "人工智能是什么？请从百度百科中查找相关信息"
})
print("回答:", response1['output'])

# 测试2：对话历史测试
print("\n--- 测试2：对话历史测试 ---")
response2 = agent_executor.invoke({
    "chat_history": [
        HumanMessage(content="Hi,我的名字是思琪"),
        AIMessage(content="你好,思琪,很高兴见到你！有什么我可以帮助你的吗？"),
    ],
    "input": "我的名字是什么？"  # 测试Agent是否能记住对话历史
})
print("回答:", response2['output'])

# 测试3：复杂问题（需要检索知识）
print("\n--- 测试3：复杂知识问题 ---")
response3 = agent_executor.invoke({
    "input": "人工智能有哪些主要应用领域？"
})
print("回答:", response3['output'])

# ==================== 系统信息展示 ====================

print("\n" + "="*50)
print("系统配置信息")
print("="*50)
print(f"大语言模型: {llm.model_name}")
print(f"向量化模型: {embeddings.model_name}")
print(f"文本块数量: {len(documents)}")
print(f"可用工具: {[tool.name for tool in tools]}")
print(f"检索配置: 返回前 {retriever.search_kwargs.get('k', '未知')} 个相关文档")

"""
扩展功能建议：

# 1. 添加更多工具
weather_tool = create_retriever_tool(another_retriever, "weather_search", "查询天气信息")
calculator_tool = Tool(name="calculator", func=calculate, description="数学计算")
tools.extend([weather_tool, calculator_tool])

# 2. 使用对话历史管理（取消注释使用）
# from langchain_community.chat_message_histories import ChatMessageHistory
# from langchain_core.runnables.history import RunnableWithMessageHistory
# 
# message_history = ChatMessageHistory()
# agent_with_chat_history = RunnableWithMessageHistory(
#     agent_executor,
#     lambda session_id: message_history,
#     input_messages_key="input",
#     history_messages_key="chat_history",
# )

# 3. 添加多个知识源
# loader1 = WebBaseLoader("url1")
# loader2 = TextLoader("local_file.txt")
# loader3 = PDFLoader("document.pdf")
# all_docs = loader1.load() + loader2.load() + loader3.load()
"""

print("\n" + "="*50)
print("学习总结")
print("="*50)
print("✓ 掌握了RAG系统的完整构建流程")
print("✓ 理解了文档加载、分割、向量化的作用")
print("✓ 学会了创建检索工具和智能代理")
print("✓ 了解了对话历史的管理方式")
print("✓ 实践了多轮对话和知识检索的结合")

print("\n下一步学习建议：")
print("1. 尝试添加更多类型的数据源（PDF、TXT、数据库等）")
print("2. 探索不同的向量化模型和检索策略")
print("3. 学习如何优化提示词模板提升代理性能")
print("4. 实践多工具协同工作的复杂代理系统")